package com.example.demo.aop;

/**
 * @Description TODO
 * @Date 2024/3/12 9:39
 * @Author ShaoJie Cen
 */

import cn.hutool.core.collection.ListUtil;
import com.example.demo.enums.LimitTypeEnum;
import com.example.demo.enums.ResultCode;
import com.example.demo.utils.IpUtil;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.reflect.MethodSignature;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * 限流处理
 */
@Aspect
@Component
public class RateLimiterAspect
{
    private static final Logger log = LoggerFactory.getLogger(RateLimiterAspect.class);



    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    @Autowired
    private RedisScript<Boolean> limitScript;



    /**
     * 限流
     * XXX 对限流要求比较高，可以使用在 Redis中对规则进行存储校验 或者使用中间件
     *
     * @param joinPoint   joinPoint
     * @param rateLimiter 限流注解
     */
    @Before(value = "@annotation(rateLimiter)")
    public void boBefore(JoinPoint joinPoint, RateLimiter rateLimiter) {
        // 1. 生成 key
        String key = getCombineKey(rateLimiter, joinPoint);
        try {
            // 2. 执行脚本返回是否限流
            Boolean flag = redisTemplate.execute(limitScript,
                    ListUtil.of(key, String.valueOf(System.currentTimeMillis())),
                    (Object[]) getRules(rateLimiter));
            // 3. 判断是否限流
            if (Boolean.TRUE.equals(flag)) {
                log.error("ip: '{}' 拦截到一个请求 RedisKey: '{}'",
                        IpUtil.getIpAddr(((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest()),
                        key);
                throw new RuntimeException(ResultCode.REQUEST_MORE_ERROR.getMessage());
            }
        } catch (RuntimeException e) {
            throw e;
        } catch (Exception e) {
            log.error("限流报错！",e);
        }
    }

    /**
     * 获取规则
     *
     * @param rateLimiter 获取其中规则信息
     * @return
     */
    private Long[] getRules(RateLimiter rateLimiter) {
        int capacity = rateLimiter.rules().length << 1;
        // 1. 构建 args
        Long[] args = new Long[rateLimiter.preventDuplicate() ? capacity + 2 : capacity];
        // 3. 记录数组元素
        int index = 0;
        // 2. 判断是否需要添加防重复提交到redis进行校验
        if (rateLimiter.preventDuplicate()) {
            RateRule preventRateRule = rateLimiter.preventDuplicateRule();
            args[index++] = preventRateRule.count();
            args[index++] = preventRateRule.timeUnit().toMillis(preventRateRule.time());
        }
        RateRule[] rules = rateLimiter.rules();
        for (RateRule rule : rules) {
            args[index++] = rule.count();
            args[index++] = rule.timeUnit().toMillis(rule.time());
        }
        return args;
    }




    /**
     * 通过 rateLimiter 和 joinPoint 拼接  prefix : ip / userId : classSimpleName - methodName
     *
     * @param rateLimiter 提供 prefix
     * @param joinPoint   提供 classSimpleName : methodName
     * @return
     */
    public String getCombineKey(RateLimiter rateLimiter, JoinPoint joinPoint) {
        StringBuffer key = new StringBuffer(rateLimiter.key());
        // 不同限流类型使用不同的前缀
        switch (rateLimiter.limitType()) {
            // XXX 可以新增通过参数指定参数进行限流
            case IP:
                key.append(IpUtil.getIpAddr(((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest())).append(":");
                break;
//            case USER_ID:
//                SysUserDetails user = SecurityUtil.getUser();
//                if (!ObjectUtils.isEmpty(user)) key.append(user.getUserId()).append(":");
//                break;
            case GLOBAL:
                break;
        }
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        key.append(targetClass.getSimpleName()).append("-").append(method.getName());
        return key.toString();
    }
}
